﻿namespace AdvancedTraining.Lesson10;

// 本题测试链接：https://www.lintcode.com/problem/top-k-frequent-words-ii/
// 以上的代码不要粘贴, 把以下的代码粘贴进java环境编辑器
// 把类名和构造方法名改成TopK, 可以直接通过
public class TopK(int k)
{
    private readonly NodeHeapComp _comp = new();

    private readonly Node[] _heap = new Node[k];

    private readonly Dictionary<Node, int> _nodeIndexMap = new();

    // 词频表   key  abc   value  (abc,7)
    private readonly Dictionary<string, Node> _strNodeMap = new();
    private readonly SortedSet<Node> _treeSet = new(new NodeTreeSetComp());
    private int _heapSize;

    public void Add(string str)
    {
        if (_heap.Length == 0) return;
        // str   找到对应节点  curNode
        Node? curNode = null;
        // 对应节点  curNode  在堆上的位置
        var preIndex = -1;
        if (!_strNodeMap.TryGetValue(str, out var value))
        {
            curNode = new Node(str, 1);
            _strNodeMap[str] = curNode;
            _nodeIndexMap[curNode] = -1;
        }
        else
        {
            curNode = value;
            // 要在time++之前，先在treeSet中删掉
            // 原因是因为一但times++，curNode在treeSet中的排序就失效了
            // 这种失效会导致整棵treeSet出现问题
            if (_treeSet.Contains(curNode)) _treeSet.Remove(curNode);
            curNode.Times++;
            preIndex = _nodeIndexMap[curNode];
        }

        if (preIndex == -1)
        {
            if (_heapSize == _heap.Length)
            {
                if (_comp.Compare(_heap[0], curNode) < 0)
                {
                    _treeSet.Remove(_heap[0]);
                    _treeSet.Add(curNode);
                    _nodeIndexMap[_heap[0]] = -1;
                    _nodeIndexMap[curNode] = 0;
                    _heap[0] = curNode;
                    Heapify(0, _heapSize);
                }
            }
            else
            {
                _treeSet.Add(curNode);
                _nodeIndexMap[curNode] = _heapSize;
                _heap[_heapSize] = curNode;
                HeapInsert(_heapSize++);
            }
        }
        else
        {
            _treeSet.Add(curNode);
            Heapify(preIndex, _heapSize);
        }
    }

    public virtual IList<string> GetTopK()
    {
        var ans = new List<string>();
        foreach (var node in _treeSet) ans.Add(node.Str);
        return ans;
    }

    /// <summary>
    /// 测试TopK高频词统计系统
    /// </summary>
    public static void Run()
    {
        Console.WriteLine("TopK高频词统计系统测试开始");

        // 测试用例1: 基本功能测试
        Console.WriteLine("测试用例1: 基本功能测试 (K=3)");
        var topK1 = new TopK(3);
        topK1.Add("i");
        topK1.Add("love");
        topK1.Add("leetcode");
        topK1.Add("i");
        topK1.Add("leetcode");
        topK1.Add("coding");
        var result1 = topK1.GetTopK();
        Console.WriteLine($"结果: [{string.Join(", ", result1)}]");
        Console.WriteLine("期望: [i, leetcode, love] 或 [i, leetcode, coding]");

        // 测试用例2: 超过K个不同单词
        Console.WriteLine("\n测试用例2: 超过K个不同单词 (K=2)");
        var topK2 = new TopK(2);
        string[] words = { "a", "b", "c", "a", "b", "a" };
        foreach (var word in words)
        {
            topK2.Add(word);
        }
        var result2 = topK2.GetTopK();
        Console.WriteLine($"结果: [{string.Join(", ", result2)}]");
        Console.WriteLine("期望: [a, b]");

        // 测试用例3: 相同频率的单词
        Console.WriteLine("\n测试用例3: 相同频率的单词 (K=3)");
        var topK3 = new TopK(3);
        string[] words3 = { "apple", "banana", "apple", "banana", "cherry" };
        foreach (var word in words3)
        {
            topK3.Add(word);
        }
        var result3 = topK3.GetTopK();
        Console.WriteLine($"结果: [{string.Join(", ", result3)}]");
        Console.WriteLine("期望: [apple, banana, cherry] (按字典序)");

        // 测试用例4: K=0的特殊情况
        Console.WriteLine("\n测试用例4: K=0的特殊情况");
        var topK4 = new TopK(0);
        topK4.Add("test");
        var result4 = topK4.GetTopK();
        Console.WriteLine($"结果: [{string.Join(", ", result4)}]");
        Console.WriteLine("期望: []");

        // 测试用例5: 动态更新测试
        Console.WriteLine("\n测试用例5: 动态更新测试 (K=2)");
        var topK5 = new TopK(2);
        topK5.Add("hello");
        topK5.Add("world");
        var result5a = topK5.GetTopK();
        Console.WriteLine($"添加hello,world后: [{string.Join(", ", result5a)}]");

        topK5.Add("hello");
        topK5.Add("hello");
        var result5b = topK5.GetTopK();
        Console.WriteLine($"再次添加hello后: [{string.Join(", ", result5b)}]");

        topK5.Add("world");
        topK5.Add("world");
        topK5.Add("world");
        var result5c = topK5.GetTopK();
        Console.WriteLine($"多次添加world后: [{string.Join(", ", result5c)}]");
        Console.WriteLine("期望: [world, hello]");

        // 测试用例6: 大量数据测试
        Console.WriteLine("\n测试用例6: 大量数据测试 (K=5)");
        var topK6 = new TopK(5);
        var random = new Random(42); // 固定种子保证结果可重现
        var testData = new string[] { "a", "b", "c", "d", "e", "f", "g", "h", "i", "j" };

        // 模拟不同频率的数据
        var frequencies = new int[] { 15, 12, 8, 6, 5, 4, 3, 2, 1, 1 };
        for (int i = 0; i < testData.Length; i++)
        {
            for (int j = 0; j < frequencies[i]; j++)
            {
                topK6.Add(testData[i]);
            }
        }

        var result6 = topK6.GetTopK();
        Console.WriteLine($"结果: [{string.Join(", ", result6)}]");
        Console.WriteLine("期望: [a, b, c, d, e]");

        Console.WriteLine("\nTopK高频词统计系统测试结束");
    }

    private void HeapInsert(int index)
    {
        while (index != 0)
        {
            var parent = (index - 1) / 2;
            if (_comp.Compare(_heap[index], _heap[parent]) < 0)
            {
                Swap(parent, index);
                index = parent;
            }
            else
            {
                break;
            }
        }
    }

    private void Heapify(int index, int heapSize)
    {
        var l = index * 2 + 1;
        var r = index * 2 + 2;
        var smallest = index;
        while (l < heapSize)
        {
            if (_comp.Compare(_heap[l], _heap[index]) < 0) smallest = l;
            if (r < heapSize && _comp.Compare(_heap[r], _heap[smallest]) < 0) smallest = r;
            if (smallest != index)
                Swap(smallest, index);
            else
                break;
            index = smallest;
            l = index * 2 + 1;
            r = index * 2 + 2;
        }
    }

    private void Swap(int index1, int index2)
    {
        _nodeIndexMap[_heap[index1]] = index2;
        _nodeIndexMap[_heap[index2]] = index1;
        (_heap[index1], _heap[index2]) = (_heap[index2], _heap[index1]);
    }

    private class Node(string s, int t)
    {
        public readonly string Str = s;
        public int Times = t;
    }

    private class NodeHeapComp : IComparer<Node>
    {
        public int Compare(Node? o1, Node? o2)
        {
            return o1!.Times != o2!.Times
                ? o1.Times - o2.Times
                : string.Compare(o2.Str, o1.Str, StringComparison.Ordinal);
        }
    }

    private class NodeTreeSetComp : IComparer<Node>
    {
        public int Compare(Node? o1, Node? o2)
        {
            return o1!.Times != o2!.Times
                ? o2.Times - o1.Times
                : string.Compare(o1.Str, o2.Str, StringComparison.Ordinal);
        }
    }
}